from config import conf
from os.path import join
import pandas as pd
import numpy as np


def choose_keys(predicts, entities):
    predicts = eval(predicts)
    entities = eval(entities)
    keys = []
    for i in range(len(predicts)):
        if predicts[i] == 1:
            keys.append(entities[i])
    return ';'.join(keys)


def reduce_rs_by_id(raw_df):
    df = raw_df[['id', 'negative', 'predict', 'entity_list']]
    df['key_entity'] = df.apply(lambda x: choose_keys(x[2], x[3]), axis=1)
    items = []
    for id, sub in df[['id', 'negative', 'key_entity']].groupby(['id']):
        negative = 1 if np.mean(sub['negative'].values) >= 0.5 else 0
        key_entity = ';'.join(filter(lambda x: x != '', sub['key_entity'].values)).strip() if negative == 1 else np.nan
        items.append((id, negative, key_entity))
    return pd.DataFrame(items, columns=['id', 'negative', 'key_entity'])


if __name__ == '__main__':
    INFERENCE_DIR = conf.get('dir', 'inference_result_dir')
    RESULT_FILE = join(INFERENCE_DIR, 'evaluation', 'BertSentiEntityscore0.940363_epoch6')

    df = pd.read_csv(RESULT_FILE)[['id', 'negative', 'predict', 'entity_list']]
    print(reduce_rs_by_id(df))
